__author__ = 'sdjkx'

# 检查指定目录里的 train 文件，保存标签数据到数据库 
import os
from libs.imageBndBox import ImageBndBox
from libs.fileInfoManager import FileInfoManager
from libs.labelFile import LabelFile
from libs.ustr import *
from libs.pascal_voc_io import PascalVocWriter
from xml.etree.ElementTree import Element, SubElement
from PyQt5.QtWidgets import QMessageBox,QFileDialog
import shutil
try:
  import xml.etree.cElementTree as ET
except ImportError:
  import xml.etree.ElementTree as ET

class ScanImageBndbox:

    def __init__(self, labelImg, mydb, userId, label_dict):
        self.labelImg = labelImg
        self.mydb = mydb
        self.userId = userId
        self.label_dict = label_dict
        self.fileInfoManager = FileInfoManager(labelImg, mydb)
        self.boxList = []
        self.rowid = 0

        self.selectSql = "select id,file_id,row_id,label_id,xmin,ymin,xmax,ymax,user_id, create_time from u_image_bndbox where file_id={0}"
        self.updateSsql = "update u_image_bndbox set label_id= {0},xmin={1},ymin={2},xmax={3},ymax={4},create_time=SYSDATE(),user_id='{5}' where id={6}"
        self.deletesql = "delete from u_image_bndbox where id in ({0})"
        self.insertSql = "insert into u_image_bndbox(file_id,row_id,label_id,xmin,ymin,xmax,ymax,create_time,user_id)" \
                         " VALUES ({0},{1},{2},{3},{4},{5},{6},SYSDATE(),{7})"

    def scanDirXml(self, userId, labelFile):
        self.labelFile = labelFile
        self.userId = userId
        dirpath = ustr(QFileDialog.getExistingDirectory(self.labelImg,
                                                        "选择XML文件路径", "",  QFileDialog.ShowDirsOnly
                                                        | QFileDialog.DontResolveSymlinks))
        if dirpath is not None and len(dirpath) > 1:
            if dirpath.endswith("train") == False:
                QMessageBox.information(self.labelImg, "","目录【" + dirpath+ "】必须是train目录")
                return

            ret = QMessageBox.question(self.labelImg, '提示', "确定要扫描目录【"+dirpath+"】对应的Annotations目录里的XML文件吗？",
                                       QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
            if ret == QMessageBox.Yes:
                xmlPath = dirpath.replace("train", "Annotations")
                if not os.path.exists(xmlPath):
                    QMessageBox.information(self.labelImg, "","目录【" + xmlPath+ "】不存在")
                    return

                self.createErrorDir(dirpath)
                self.scanLabel(dirpath, xmlPath)
                return

    def queryFileBndbox(self, fileId):
        sql = self.selectSql.format(fileId)
        # 执行SQL语句, 获取所有记录列表
        results = self.mydb.executeSql(sql)

        self.boxList.clear()
        for row in results:
            box = ImageBndBox(row[0], row[1], row[2], row[3],row[4], row[5],row[6], row[7], row[8], row[9])
            self.boxList.append(box)

    def RefreshBndbox(self,fileId,userId, shapes):
        self.userId = userId
        self.rowid = 0
        for shape in shapes:
            shapeRowid = 0
            if "rowid" in shape:
                shapeRowid = int(shape["rowid"])
            else:
                shape["rowid"] = "0"
            if shapeRowid > 0:
                self.rowid = shapeRowid

        for shape in shapes:
            rowIdText = shape["rowid"]
            if rowIdText == "0":
                self.rowid += 1
                shape["rowid"] = self.rowid

        for shape in shapes:
            points = shape['points']
            label = shape['label']
            # Add Chris
            difficult = int(shape['difficult'])
            if "rowid" in shape:
                self.rowid = int(shape["rowid"])
            else:
                self.rowid = 0

            labelId = self.label_dict[label].get_id()
            bndbox = LabelFile.convertPoints2BndBox(points)
            xmin = bndbox[0]
            ymin = bndbox[1]
            xmax = bndbox[2]
            ymax = bndbox[3]
            # addBndBox(bndbox[0], bndbox[1], bndbox[2], bndbox[3], label,rowid, difficult)
            if self.rowid == 0:
                box = ImageBndBox(0, fileId, self.rowid, labelId, xmin,ymin, xmax, ymax,self.userId)
                self.boxList.append(box)
            else:
                box = self.getLableByRowId(self.rowid)
                if box is None:
                    box = ImageBndBox(0, fileId, self.rowid, labelId, xmin,ymin, xmax, ymax,self.userId)
                    self.boxList.append(box)
                elif box.label_id != labelId or box.xmin != xmin or box.ymin != ymin or box.xmax != xmax or box.ymax != ymax:
                    #  先删除
                    self.boxList.remove(box)
                    box.setLableId(labelId)
                    box.setXMin(xmin)
                    box.setYMin(ymin)
                    box.setXMax(xmax)
                    box.setYMax(ymax)

        self.updateFileBndbox()

    def getBoxById(self,boxId):
        for box in self.boxList:
            if(box.id == boxId):
                return box

        return None

    def delBndbox(self, shape):
        rowid = int(shape.rowid)
        box = self.getLableByRowId(rowid)
        if box is not None:
            box.delBox()

    def updateFileBndbox(self):
        deleteIdList = ""

        for box in self.boxList:
            if box.id == 0:
                self.insertBndbox(box)

            if box.deleted:
                if deleteIdList == "":
                    deleteIdList += str(box.id)
                else:
                    deleteIdList += "," + str(box.id)
            elif box.changed:
                sql = self.updateSsql.format(box.label_id, box.xmin, box.ymin, box.xmax, box.ymax, self.userId, box.id)
                print(sql)
                self.mydb.executeCommit(sql)

        if len(deleteIdList) > 0:
            sql = self.deletesql.format(deleteIdList)
            print(sql)
            self.mydb.executeCommit(sql)

    def insertBndbox(self, box):
        sql = self.insertSql.format(box.file_id, box.row_id,box.label_id,box.xmin,box.ymin,box.xmax,box.ymax,self.userId)
        # print(sql)
        self.mydb.executeCommit(sql)

    def scanLabel(self, jpgFileDir, xmlFileDir):
        '''
        扫描目录里的XML文件，保存到数据里去
        :param jpgFilePath: 图片目录
        :param xmlFilePath: XML文件目录
        :return:
        '''
        jpgFileList = os.listdir(jpgFileDir)
        for jpgFileName in jpgFileList:
            jpgFilePath = os.path.join(jpgFileDir, jpgFileName)
            xmlFilePath = os.path.join(xmlFileDir,jpgFileName.split(".")[0] + ".xml")
            print("检测文件：" + xmlFilePath)

            # XML文件不存在，跳过
            if not os.path.exists(xmlFilePath):
                continue

            fileId = self.getFileId(jpgFileName)
            self.queryFileBndbox(fileId)

            self.rowid = 0
            tree = ET.parse(xmlFilePath)     # 打开xml文档
            root = tree.getroot()            #获得root节点
            change = False
            objectList = root.findall('object')

            #  获取最大的rowid
            for object in objectList:
                shapeRowid = 0
                rowIdObject = object.find("rowid")
                if rowIdObject is not None:
                    shapeRowid = int(rowIdObject.text)
                else:
                    rowIdObject = SubElement(object, "rowid")
                    rowIdObject.text = "0"
                if shapeRowid > 0:
                    self.rowid = shapeRowid

            # 查找没有设置rowid 的对象，设置一个rowid
            for object in objectList:
                rowIdObject = object.find("rowid")
                if rowIdObject.text == "0":
                    self.rowid += 1
                    rowIdObject.text = self.rowid
                    change = True

            labelError = False
            for object in objectList: #找到root节点下的所有object节点
                labelError = not self.checkObject(self.rowid, fileId, object)

            if change:
                # tree.write(xmlFilePath,encoding = 'utf8')          # 保存文件
                self.save(root,jpgFilePath, xmlFilePath)
            if labelError:
                self.copyFileToError(jpgFilePath, xmlFilePath)

        QMessageBox.information(self.labelImg, "","目录【" + xmlFileDir + "】数据处理完成")

    def save(self, root,jpgFilePath, xmlFilePath):
        imgFolderPath = os.path.dirname(jpgFilePath)
        imgFolderName = os.path.split(imgFolderPath)[-1]
        imgFileName = os.path.basename(jpgFilePath)
        size = root.find("size")
        height = int(size.find("height").text)
        width = int(size.find("width").text)
        imageShape = [height, width, 3]
        writer = PascalVocWriter(imgFolderName, imgFileName,
                                 imageShape, localImgPath=jpgFilePath)
        writer.verified = self.labelFile.verified

        objectList = root.findall('object')
        for object in objectList:
            label = object.find("name").text
            difficult = int(object.find('difficult').text)
            rowid = int(object.find("rowid").text)
            bndbox = object.find('bndbox')      # 子节点下属性bndbox的值
            xmin = int(bndbox.find('xmin').text)
            ymin = int(bndbox.find('ymin').text)
            xmax = int(bndbox.find('xmax').text)
            ymax = int(bndbox.find('ymax').text)
            writer.addBndBox(xmin,ymin, xmax, ymax, label, rowid, difficult)

        writer.save(targetFile=xmlFilePath)

    def getLableByRowId(self, rowid):
        '''
        从标签字典中获取标签ID
        :param rowid: 行号
        :return:
        '''
        for box in self.boxList:
            if box.row_id == rowid:
                return box
        return None

    def getFileId(self, jpgFileName):
        '''
        获取文件ID
        :param jpgFileName: 文件名称
        :return: 文件ID
        '''
        userId, fileId, fileStatus = self.fileInfoManager.queryFile(jpgFileName)
        if fileId == 0:
            self.fileInfoManager.insertFile("100007", jpgFileName, 1)
            userId, fileId, fileStatus = self.fileInfoManager.queryFile(jpgFileName)

        return fileId

    def checkObject(self,rowid,fileId, object):
        '''
        检测每个XML文件的内容里，各个框是否有保存到数据库里，如果没有就保存
        如果
        :param rowid: 行叫
        :param fileId: 文件ID
        :param object: 标签框对象
        :return: 如果标签rowid是0是返回 rowid,如果不是0，则返回行号
        '''
        label_name = object.find('name').text   # 子节点下节点name的值
        if label_name not in self.label_dict:
            return False
        lableId = self.label_dict[label_name].get_id()
        rowidObjct = object.find("rowid")
        rowidText = rowidObjct.text

        bndbox = object.find('bndbox')      # 子节点下属性bndbox的值
        xmin = float(bndbox.find('xmin').text)
        ymin = float(bndbox.find('ymin').text)
        xmax = float(bndbox.find('xmax').text)
        ymax = float(bndbox.find('ymax').text)
        
        rowid = int(rowidText)
        box = self.getLableByRowId(rowid)
        if box is None:
            box = ImageBndBox(0, fileId, rowid, lableId, xmin,ymin, xmax, ymax,self.userId)
            self.insertBndbox(box)
        return True

    def copyFileToError(self, jpgFilePath, xmlFilePath):
        imgFileName = os.path.basename(jpgFilePath)
        xmlFileName = os.path.basename(xmlFilePath)
        if not os.path.exists(self.trainPath):
            os.mkdir(self.trainPath)
        if not os.path.exists(self.annotationsPath):
            os.mkdir(self.annotationsPath)

        jpgObjFilePath = os.path.join(self.trainPath, imgFileName)
        xmlObjFilePath = os.path.join(self.annotationsPath, xmlFileName)
        shutil.copy(jpgFilePath, jpgObjFilePath)
        shutil.copy(xmlFilePath, xmlObjFilePath)

    def createErrorDir(self, imgFolderPath):
        # imgFolderPath = os.path.dirname(jpgFilePath)
        dirList = str(imgFolderPath).split('/')
        dateName = dirList[-2]
        rootPath = ""
        dirTemp = dirList[:-2]
        for dir in dirTemp:
            rootPath += dir + "/"

        rootPath += "error/" + dateName
        self.trainPath = rootPath + "/train"
        self.annotationsPath = rootPath + "/AnnotationsPath"

